# Descriptive maps
library(pacman)
p_load(
  data.table, here, ggplot2, fixest, magrittr, dplyr, ggExtra, 
  haven, purrr, sf, tigris, tidyverse, janitor, 
  latex2exp, scales, RColorBrewer
)
options(tigris_use_cache = TRUE)
WIDTH = 8 #inches
theme_set(
  theme_void(base_size = 14) +
    theme(
      legend.position = 'bottom',
      plot.margin = margin(0,0,0,0),
      legend.key.width = unit(WIDTH/6, 'cm'),
      legend.text = element_text(size = 18)
    )
)

# Loading shape files 
states_sf = 
  states(year = 2020, cb = TRUE) |>
  filter( # Limiting to continental US
    !(STATEFP %in% c("02","15","60","66","69","72","78"))
  ) |>
  clean_names() |> 
  st_transform(crs = 2163)
tract_sf =
  map_dfr(
    states_sf$statefp,
    \(st){tracts(cb = TRUE, year = 2015, state = st)}
  ) |> 
  clean_names() |>
  st_transform(crs = 2163)

# Installations per capita from deepsolar data --------------------------------
  # First loading deepsolar data 
  deepsolar_dt = 
    fread(
      here("Data/deep-solar/deepsolar_tract.csv")
    )[,':='(
      fips = str_pad(fips, width = 11, side = "left", pad = "0"),
      solar_systems_per_cap = ifelse(
        population == 0, 0, 
        solar_system_count/population
      ),
      solar_systems_residential_per_cap = ifelse(
        population == 0, 0, 
        solar_system_count_residential/population
      ),
      pr_rad = avg_electricity_retail_rate*daily_solar_radiation
    )]
  # Calculating quantiles of solar system counts 
  solar_systems_per_cap_quant = 
    quantile(
      deepsolar_dt[solar_systems_per_cap>0]$solar_systems_per_cap, 
      probs = seq(0,1,by = 0.01)
    )
  solar_systems_residential_per_cap_quant = 
    quantile(
      deepsolar_dt[solar_systems_residential_per_cap>0]$solar_systems_residential_per_cap, 
      probs = seq(0,1,by = 0.01)
    )
  # Saving these for result plots 
  install_quant_dt = 
    data.table(
      quant = names(solar_systems_per_cap_quant),
      probs = seq(0,1,by = 0.01),
      value = solar_systems_per_cap_quant
    )
  fwrite(
    install_quant_dt, 
    here('Data/deep-solar/installs-per-cap-quantiles.csv')
  )
  # Adding bins to the data 
  deepsolar_dt[,':='(
    solar_systems_per_cap_bin = findInterval(
      solar_systems_per_cap, 
      solar_systems_per_cap_quant,
      left.open = TRUE
    ),
    solar_systems_residential_per_cap_bin = findInterval(
      solar_systems_residential_per_cap, 
      solar_systems_residential_per_cap_quant,
      left.open = TRUE
    )
  )]
  # Merging with tract shape
  deepsolar_sf = 
    left_join(
      tract_sf,
      deepsolar_dt,
      by = c("geoid"="fips")
    ) 
  # Creating map
  solar_systems_per_capita_p = 
    ggplot() +
    geom_sf(
      data = deepsolar_sf,#|> filter(statefp == '06'), 
      aes(fill = solar_systems_per_cap_bin, color = solar_systems_per_cap_bin),
      lwd = 0.01
    ) +
    geom_sf(data = states_sf, fill = NA, color = 'black') + 
    scale_color_distiller(
      name = '', 
      type = 'seq',
      palette = 'YlGnBu',
      direction = 1,
      aesthetics = c('color','fill'),
      breaks = c(1,33,67,95),
      labels = c(
        round(1000*solar_systems_per_cap_quant['1%'], digits = 1),
        round(1000*solar_systems_per_cap_quant['33%'], digits = 1),
        round(1000*solar_systems_per_cap_quant['67%'], digits = 1),
        round(1000*solar_systems_per_cap_quant['95%'], digits = 0)
      ),
      guide = guide_colorbar(ticks.colour = "white", ticks.linewidth = 1.25)
    )
  ggsave(
    plot = solar_systems_per_capita_p,
    filename = here('figures/maps/descriptive/solar-systems-per-capita.jpeg'),
    bg = 'white', 
    width = WIDTH, height = WIDTH/1.4, units = 'in'
  )

# Now map of daily solar radiation --------------------------------------------
  # First we have to fill in missing values 
  geoid_solar_radiation_na = filter(deepsolar_sf, is.na(daily_solar_radiation))$geoid
  # Joining to tracts bordering 
  border_missing_tract = 
    st_join(
      tract_sf |> filter(geoid %in% geoid_solar_radiation_na),
      deepsolar_sf |> filter(!(geoid %in% geoid_solar_radiation_na)) |> select(daily_solar_radiation),
      join = st_touches
    )
  # Taking simple average over all bordering tracts
  imputed_solar_radiation_dt = 
    data.table(border_missing_tract)[,.(
      daily_solar_radiation_imp = mean(daily_solar_radiation, na.rm = TRUE)),
      keyby = geoid
  ]
  # Adding back into the data 
  deepsolar_sf = 
    left_join(
      deepsolar_sf, 
      imputed_solar_radiation_dt, 
      by = 'geoid'
    ) |>
    mutate(
      daily_solar_radiation = ifelse(
        is.na(daily_solar_radiation), daily_solar_radiation_imp,
        daily_solar_radiation
      )
    ) 
  # Creating map
  daily_solar_radiation_p = 
    ggplot() +
    geom_sf(
      data = deepsolar_sf ,#|> filter(statefp == '01'), 
      aes(fill = daily_solar_radiation, color = daily_solar_radiation),
      lwd = 0.01
    ) +
    scale_color_viridis_c(
      name = '',# TeX(r"(Daily solar radiation $\left(kWh/m^2/day\right)$)"), 
      option = 'magma',
      aesthetics = c('color','fill'),
      guide = guide_colorbar(ticks.colour = "white", ticks.linewidth = 1.25)
    ) 
  ggsave(
    plot = daily_solar_radiation_p,
    filename = here('figures/maps/descriptive/daily-solar-radiation.jpeg'),
    bg = 'white', 
    width = WIDTH, height = WIDTH/1.4, units = 'in'
  )

# State level maps ------------------------------------------------------------
  # Reading in expected subsidy data 
  state_subs = fread(
    here('Data/CleanedData/StateAveSubs.csv')
  )[,state_fips := str_pad(state_fips, 2, 'left','0')]
  # Creating rankings of sub_ben, tot_ben, price
  state_subs[,':='(
    sub_ben_rank = frank(sub_ben),
    tot_ben_rank = frank(tot_ben),
    price_rank = frank(price)
  )]
  # Merging 
  state_sub_sf = 
    left_join(
      states_sf, 
      state_subs, 
      by = c('statefp'='state_fips')
    )
  # First plotting expected subsidy for a 15 panel system
  expected_subsidy_state_p = 
    ggplot() + 
    geom_sf(data = state_sub_sf, aes(fill = sub_ben_rank), color= 'black') +
    scale_fill_distiller(
      name = '',#'Expected subsidy ($1000)', 
      type = 'seq',
      palette = 'YlGnBu',
      direction = 1,
      breaks = c(1,16,32,48),
      labels = dollar(c(
        round(1e-3*state_subs[sub_ben_rank == 1]$sub_ben[1], digits = 1),
        round(1e-3*state_subs[sub_ben_rank == 16]$sub_ben[1], digits = 1),
        round(1e-3*state_subs[sub_ben_rank == 32]$sub_ben[1], digits = 1),
        round(1e-3*state_subs[sub_ben_rank == 48]$sub_ben[1], digits = 0)
      ), largest_with_cents = 0, suffix = 'K'),
      guide = guide_colorbar(ticks.colour = "white", ticks.linewidth = 1.25)
    )
  ggsave(
    plot = expected_subsidy_state_p,
    filename = here('figures/maps/descriptive/expected-subsidy.jpeg'),
    bg = 'white', 
    width = WIDTH, height = WIDTH/1.4, units = 'in'
  )
  # Expected monetary benefit for a 15 panel system
  expected_monetary_benefit_state_p = 
    ggplot() + 
    geom_sf(
      data = state_sub_sf, 
      aes(fill = tot_ben_rank),#as.factor(findInterval(tot_ben, c(-5000,0,5000)))), 
      color= 'black'
    ) + 
    scale_fill_distiller(
      name = '',#'Expected monetary benefit ($1000)', 
      type = 'seq',
      palette = 'YlGnBu',
      direction = 1,
      breaks = c(1,16,32,48),
      labels = dollar(c(
        round(1e-3*state_subs[tot_ben_rank == 1]$tot_ben[1], digits = 1),
        round(1e-3*state_subs[tot_ben_rank == 16]$tot_ben[1], digits = 1),
        round(1e-3*state_subs[tot_ben_rank == 32]$tot_ben[1], digits = 1),
        round(1e-3*state_subs[tot_ben_rank == 48]$tot_ben[1], digits = 0)
      ), largest_with_cents = 0, suffix = 'K'),
      guide = guide_colorbar(ticks.colour = "white", ticks.linewidth = 1.25)
    )
  ggsave(
    plot = expected_monetary_benefit_state_p,
    filename = here('figures/maps/descriptive/expected-monetary-benefit.jpeg'),
    bg = 'white', 
    width = WIDTH, height = WIDTH/1.4, units = 'in'
  )
  # State electricity prices
  elec_prices_state_p = 
    ggplot() + 
    geom_sf(data = state_sub_sf, aes(fill = price_rank), color= 'black') + 
    scale_fill_distiller(
      name = '',#'Electricity price ($/kWh)', 
      type = 'seq',
      palette = 'YlGnBu',
      direction = 1,
      breaks = c(1,16,32,48),
      labels = dollar(c(
        round(state_subs[price_rank ==  1]$price[1], digits = 2),
        round(state_subs[price_rank == 16]$price[1], digits = 2),
        round(state_subs[price_rank == 32]$price[1], digits = 2),
        round(state_subs[price_rank == 48]$price[1], digits = 2)
      )),
      guide = guide_colorbar(ticks.colour = "white", ticks.linewidth = 1.25)
    ) 
  ggsave(
    plot = elec_prices_state_p,
    filename = here('figures/maps/descriptive/electricity-prices.jpeg'),
    bg = 'white', 
    width = WIDTH, height = WIDTH/1.4, units = 'in'
  )
  # Three types of subsidies broken out 
  unit_ben_p = 
    ggplot() + 
    geom_sf(data = state_sub_sf, aes(fill = unit_ben/1e3), color= 'black') +
    scale_fill_distiller(
      name = '', 
      type = 'seq',
      palette = 'YlGnBu',
      direction = 1,
      # breaks = c(4,16,32,45),
      labels = label_dollar(largest_with_cents = 0, suffix = 'K'),
      guide = guide_colorbar(ticks.colour = "white", ticks.linewidth = 1.25)
    )
    unit_ben_p
  ggsave(
    plot = unit_ben_p,
    filename = here('figures/maps/descriptive/unit-subsidy.jpeg'),
    bg = 'white', 
    width = WIDTH, height = WIDTH/1.4, units = 'in'
  )
  cost_ben_p = 
    ggplot() + 
    geom_sf(data = state_sub_sf, aes(fill = cost_ben/1e3), color= 'black') +
    scale_fill_distiller(
      name = '', 
      type = 'seq',
      palette = 'YlGnBu',
      direction = 1,
      # breaks = c(4,16,32,45),
       labels = label_dollar(largest_with_cents = 0, suffix = 'K'),
      guide = guide_colorbar(ticks.colour = "white", ticks.linewidth = 1.25)
    )
    cost_ben_p
  ggsave(
    plot = cost_ben_p,
    filename = here('figures/maps/descriptive/cost-subsidy.jpeg'),
    bg = 'white', 
    width = WIDTH, height = WIDTH/1.4, units = 'in'
  )
  kwh_ben_p = 
    ggplot() + 
    geom_sf(data = state_sub_sf, aes(fill = kwh_ben/1e3), color= 'black') +
    scale_fill_distiller(
      name = '', 
      type = 'seq',
      palette = 'YlGnBu',
      direction = 1,
       breaks = c(0,4,8,12),
       labels = label_dollar(largest_with_cents = 0, suffix = 'K'),
      guide = guide_colorbar(ticks.colour = "white", ticks.linewidth = 1.25)
    )
    kwh_ben_p
  ggsave(
    plot = kwh_ben_p,
    filename = here('figures/maps/descriptive/kwh-subsidy.jpeg'),
    bg = 'white', 
    width = WIDTH, height = WIDTH/1.4, units = 'in'
  )
    